/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd
 * [graph-engine] is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/


#include "parser/common/op_registration_tbe.h"
#include <map>
#include <memory>
#include <string>
#include "parser/common/acl_graph_parser_util.h"
#include "common/op_map.h"
#include "common/util.h"
#include "framework/common/debug/ge_log.h"
#include "graph/utils/type_utils_inner.h"
#include "parser/common/op_parser_factory.h"
#include "parser/tensorflow/tensorflow_custom_parser_adapter.h"
#include "parser/tensorflow/tensorflow_fusion_custom_parser_adapter.h"

namespace ge {
namespace {
std::string GetOmOptype(const OpRegistrationData &reg_data) {
  AscendString om_op_type;
  (void)reg_data.GetOmOptype(om_op_type);
  return om_op_type.GetString() == nullptr ? "" : std::string(om_op_type.GetString());
}
}
using PARSER_CREATOR_FN = std::function<std::shared_ptr<OpParser>(void)>;

FMK_FUNC_HOST_VISIBILITY OpRegistrationTbe *OpRegistrationTbe::Instance() {
  static OpRegistrationTbe instance;
  return &instance;
}

bool OpRegistrationTbe::Finalize(const OpRegistrationData &reg_data, bool is_train) {
  static std::map<domi::FrameworkType, std::map<std::string, std::string> *> op_map = {{domi::CAFFE, &caffe_op_map}};
  if (is_train) {
    op_map[domi::TENSORFLOW] = &tensorflow_train_op_map;
  } else {
    op_map[domi::TENSORFLOW] = &tensorflow_op_map;
  }

  if (op_map.find(reg_data.GetFrameworkType()) != op_map.end()) {
    std::map<std::string, std::string> *fmk_op_map = op_map[reg_data.GetFrameworkType()];
    std::set<AscendString> ori_optype_set;
    (void)reg_data.GetOriginOpTypeSet(ori_optype_set);
    for (auto &tmp : ori_optype_set) {
      if (tmp.GetString() == nullptr) {
        continue;
      }
      if ((*fmk_op_map).find(tmp.GetString()) != (*fmk_op_map).end()) {
        GELOGW("Op type does not need to be changed, om_optype:%s, orignal type:%s.",
               (*fmk_op_map)[tmp.GetString()].c_str(), tmp.GetString());
        continue;
      } else {
        (*fmk_op_map)[tmp.GetString()] = GetOmOptype(reg_data);
        GELOGD("First register in parser initialize, original type: %s, om_optype: %s, imply type: %s.",
               tmp.GetString(), GetOmOptype(reg_data).c_str(),
               TypeUtilsInner::ImplyTypeToSerialString(reg_data.GetImplyType()).c_str());
      }
    }
  }

  bool ret = RegisterParser(reg_data);
  return ret;
}

bool OpRegistrationTbe::RegisterParser(const OpRegistrationData &reg_data) const {
  if (reg_data.GetFrameworkType() == domi::TENSORFLOW) {
    std::shared_ptr<OpParserFactory> factory = OpParserFactory::Instance(domi::TENSORFLOW);
    if (factory == nullptr) {
      REPORT_CALL_ERROR("E19999", "Get OpParserFactory failed.");
      GELOGE(INTERNAL_ERROR, "[Get][OpParserFactory] for tf failed.");
      return false;
    }
    if (reg_data.GetParseParamFn() != nullptr || reg_data.GetParseParamByOperatorFn() != nullptr) {
      bool is_registed = factory->OpParserIsRegistered(GetOmOptype(reg_data));
      if (is_registed) {
        GELOGW("Parse param func has already register for op:%s.", GetOmOptype(reg_data).c_str());
        return false;
      }
      std::shared_ptr<TensorFlowCustomParserAdapter> tf_parser_adapter =
          ge::parser::MakeShared<TensorFlowCustomParserAdapter>();
      if (tf_parser_adapter == nullptr) {
        REPORT_CALL_ERROR("E19999", "Create TensorFlowCustomParserAdapter failed.");
        GELOGE(PARAM_INVALID, "[Create][TensorFlowCustomParserAdapter] failed.");
        return false;
      }
      OpParserRegisterar registerar __attribute__((unused)) = OpParserRegisterar(
          domi::TENSORFLOW, GetOmOptype(reg_data), [tf_parser_adapter]() -> std::shared_ptr<OpParser>
          { return tf_parser_adapter; });
    }
    if (reg_data.GetFusionParseParamFn() != nullptr || reg_data.GetFusionParseParamByOpFn() != nullptr) {
      bool is_registed = factory->OpParserIsRegistered(GetOmOptype(reg_data), true);
      if (is_registed) {
        GELOGW("Parse param func has already register for fusion op:%s.", GetOmOptype(reg_data).c_str());
        return false;
      }
      GELOGI("Register fusion custom op parser: %s", GetOmOptype(reg_data).c_str());
      std::shared_ptr<TensorFlowFusionCustomParserAdapter> tf_fusion_parser_adapter =
          ge::parser::MakeShared<TensorFlowFusionCustomParserAdapter>();
      if (tf_fusion_parser_adapter == nullptr) {
        REPORT_CALL_ERROR("E19999", "Create TensorFlowFusionCustomParserAdapter failed.");
        GELOGE(PARAM_INVALID, "[Create][TensorFlowFusionCustomParserAdapter] failed.");
        return false;
      }
      OpParserRegisterar registerar __attribute__((unused)) = OpParserRegisterar(
          domi::TENSORFLOW, GetOmOptype(reg_data),
          [tf_fusion_parser_adapter]() -> std::shared_ptr<OpParser> { return tf_fusion_parser_adapter; }, true);
    }
  } else {
    std::shared_ptr<OpParserFactory> factory = OpParserFactory::Instance(reg_data.GetFrameworkType());
    if (factory == nullptr) {
      REPORT_CALL_ERROR("E19999", "Get OpParserFactory for %s failed.",
                        TypeUtilsInner::FmkTypeToSerialString(reg_data.GetFrameworkType()).c_str());
      GELOGE(INTERNAL_ERROR, "[Get][OpParserFactory] for %s failed.",
             TypeUtilsInner::FmkTypeToSerialString(reg_data.GetFrameworkType()).c_str());
      return false;
    }
    bool is_registed = factory->OpParserIsRegistered(GetOmOptype(reg_data));
    if (is_registed) {
      GELOGW("Parse param func has already register for op:%s.", GetOmOptype(reg_data).c_str());
      return false;
    }

    PARSER_CREATOR_FN func = CustomParserAdapterRegistry::Instance()->GetCreateFunc(reg_data.GetFrameworkType());
    if (func == nullptr) {
      REPORT_CALL_ERROR("E19999", "Get custom parser adapter failed for fmk type %s.",
                        TypeUtilsInner::FmkTypeToSerialString(reg_data.GetFrameworkType()).c_str());
      GELOGE(INTERNAL_ERROR, "[Get][CustomParserAdapter] failed for fmk type %s.",
             TypeUtilsInner::FmkTypeToSerialString(reg_data.GetFrameworkType()).c_str());
      return false;
    }
    OpParserFactory::Instance(reg_data.GetFrameworkType())->RegisterCreator(GetOmOptype(reg_data), func);
    GELOGD("Register custom parser adapter for op %s of fmk type %s success.", GetOmOptype(reg_data).c_str(),
           TypeUtilsInner::FmkTypeToSerialString(reg_data.GetFrameworkType()).c_str());
  }
  return true;
}
}  // namespace ge
